# %% 神经网络
# 定义一个卷积神经网络
# 网络结构为：其实就是LeNet
# 卷积 池化 卷积 池化 全连接 全连接 高斯测试
import torch
import torch.nn as nn
import torch.nn.functional as F


# %%
# 定义网络，一定是继承至nn.Module
class Net(nn.Module):
    # init需要实例化网络层模块
    def __init__(self):
        # 初始化超类
        super(Net, self).__init__()
        # in:1 out:6 kernel_size:3x3
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(16 * 6 * 6, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    # forward定义前向传播过程,也就是网络结构
    def forward(self, x):
        # 第一层卷积+relu激活后最大池化
        # 用2x2的窗口进行最大池化
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # 第二层卷积+relu后进行最大池化
        # 如果形状是正方形那么你可以只指定一个数字，也就是说2就代表了2,2
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        # 进行展平处理，-1代表这里的大小没有指定，可以根据后面进行推断
        x = x.view(-1, self.num_flat_features(x))
        # 第3层接全连接后过relu
        x = F.relu(self.fc1(x))
        # 第4层接全连接后过relu
        x = F.relu(self.fc2(x))
        # 第5层全连接后输出
        x = self.fc3(x)
        # 最后记得把结果output return回去
        return x

    # 前面的forward用到了这个参数：返回特征的打平形状
    # size = (a,b,c)也就是说num_feature =1*b*c
    def num_flat_features(self, x):
        # 设样本轴为1
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

    # 其实上面的可以进行简写
    # tensor.prod()表示矩阵连乘
    # def num_flat_features(self, x):
    #     return x.size()[1:].prod()


# 实例化网络
net = Net()
# 打印网络
print(net)
# %%
# 看看参数长什么样子吧
# 注意 net.parameters()返回的是一个迭代器
# 先转换成列表再来查看
params = list(net.parameters())
# 打印一下形状呢
# 迭代器长度
print(len(params))
# 第1个的形状
print(params[0].size())
# %%
# 尝试进行32x32输入 LeNet(input:32x32)
# 当然这里的输入和输出都是无意义的random
input_ = torch.randn(1, 1, 32, 32)
out = net(input_)
print(out)
# %%
# 把参数和反向传播梯度缓冲区归零，以便继续下面的操作
net.zero_grad()
out.backward(torch.randn(1, 10))
# %%
# torch.nn仅支持小批量
# nn.Conv2d输入为n c h w 另外tensorflow的应该是n h w c
# 前向传播得out
output = net(input_)
# 设定目标，也就是标签
target = torch.randn(10)
target = target.view(1, -1)
# 定义损失函数
criterion = nn.MSELoss()
# 计算损失，并打印
loss = criterion(output, target)
print(loss)

# %%
# 反向传播
# 注意，反向传播前，要将梯度清零
net.zero_grad()
# 先看下偏置的梯度
print(net.conv1.bias.grad)
# 进行反向传播
loss.backward()
# 再看下梯度
print(net.conv1.bias.grad)
# %%
# 使用优化器,更新权重
# 随机梯度下降（SGD）
learning_rate = 0.01
# 就是 参数=参数-(损失对参数的梯度*学习率)
for f in net.parameters():
    f.data.sub_(f.grad.data * learning_rate)
#%%
# 先打印一下第一个参数
param = next(iter(net.parameters()))
print(param[0])
# %%
# 这里有更多的优化器SGD Adam RMSProp等等
import torch.optim as optim

optimizer = optim.SGD(net.parameters(), lr=0.01)
# 先清空梯度缓冲区
optimizer.zero_grad()
# 前向传播
output = net(input_)
# 计算loss
loss = criterion(output, target)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()

#%%
# 再打印一下第一个参数，看看更新了吗
param = next(iter(net.parameters()))
print(param[0])
# 确实更新了一点点